import baostock as bs
import pandas as pd
import datetime
import string
import csv
import os
import mysql

# 首次运行,获取所有股票的股票代码信息
# #update_stk_list(TODATE)
                                                  
#下载周线
#download_data(stk_list,frequency = 'w')


FROMDATE = '2022-10-30'
TODATE = '2022-11-1'

#获取指定日期的期指股票数据
def update_stk_list(date = None):
    stock_rs = bs.query_all_stock(date)
    stock_df = stock_rs.get_data()
    #stock_df.to_csv('./stk_data/all_list.csv', encoding = 'gbk', index = False)
    stock_df.drop(stock_df[stock_df.code < 'sh.600000'].index, inplace = True)
    stock_df.drop(stock_df[stock_df.code > 'sz.300001'].index, inplace = True)
    stock_df = stock_df['code']
    stock_df.to_csv('./stk_data/stk_list.csv', encoding = 'gbk', index = False)
    return stock_df.tolist()

def load_stk_list():
    df = pd.read_csv('./stk_data/stk_list.csv')
    return df['code'].tolist()

#下载股票数据
# todate = datetime.date.today(), 
def download_data(code, fromdate = '2022-10-10', todate = '2022-10-20', datas = 'date,open,high,low,close,volume,amount,turn,pctChg', 
                  frequency = 'd', adjustflag = '2'):
         print("Downloading :" + code) 
         k_rs = bs.query_history_k_data_plus(code, datas, start_date = fromdate, end_date = todate,frequency = frequency, adjustflag = adjustflag)
         
         out_df = k_rs.get_data()           
         #保存csv
         #datapath = './stk_data/' + frequency + '/' + code + '.csv'
         #out_df.to_csv(datapath, encoding = 'gbk', index = True)        
         return out_df



# #下载指定股票的日线数据
# def download_code_day_data(code,day):
#      out_df= download_data(code,frequency='d',fromdate=day,todate=day)
#      stockdata={}    
#      if not out_df.empty:
#         stockdata['date'] = out_df.iloc[0,0]
#         stockdata['open']= out_df.iloc[0,1]
#         stockdata['high']= out_df.iloc[0,2]
#         stockdata['low']= out_df.iloc[0,3]
#         stockdata['close']= out_df.iloc[0,4]
#         stockdata['volume']= out_df.iloc[0,5]
#         stockdata['amount']= out_df.iloc[0,6]
#         stockdata['turn']= out_df.iloc[0,7]
#         stockdata['pctChg']= out_df.iloc[0,8]
    

#下载并创建所有股票近一年的数据库
def download_create_all_database():
    #登录
    bs.login()
    # #连接mysql
    db=mysql.connect()
    stk_list = load_stk_list()
    
    #stk_list = {'sh.600000'}
    #创建每个股票的数据表
    for code in stk_list:
        download_create_database(code,db)
    #关闭mysql
    mysql.close(db)
    #登出
    bs.logout()
                

def download_create_database(code,db):
    """下载股票近一年数据并且创建数据库  
    参数:
        code (string): 股票代码
        db(class): pymysql数据库
    """    
    begin = datetime.date.today()-datetime.timedelta(days=365)
    end = datetime.date.today()
    download_update_stock_day(code,db,begin,end)
    

def download_update_all_stock_day(begin=datetime.date.today(),end=datetime.date.today()):            
    """下载并更新所有股票指定日期的数据
    参数:
        begin (string, optional): 开始日期. Defaults to datetime.date.today().
        end (string, optional): 结束日期. Defaults to datetime.date.today().
    Returns:
        bool:返回是否成功
    """    
    try:
        #登录
        bs.login()
        #连接mysql
        db=mysql.connect()
        stk_list = load_stk_list()
    
        #stk_list = {'sh.688776'}
        #创建每个股票的数据表
        for code in stk_list:
            download_update_stock_day(code,db,begin,end)
            
            
        #关闭mysql
        mysql.close(db)
        #登出
        bs.logout()
        return True
    except Exception as e:
        print(e)
            
#下载并且更新指定股票指定日期的数据
def download_update_stock_day(code,db,datebegin=datetime.date.today(),dateend=datetime.date.today()):
    begin = str(datebegin)
    end = str(dateend)
    df =download_data(code,begin,end)
    #获取长度
    lenth = len(df)
    
    #数据返回为空
    if df.empty:
        print("download_data%s is empty"%code)
        return None
    
    # #创建表结构
    code_name = code.replace('.','')
    mysql.create_table(db,code_name)    
    
    stockdata={}    
    for i in range(0,lenth):
        #获取每日数据
        stockdata['date'] = df.iloc[i,0]
        stockdata['open']= df.iloc[i,1]
        stockdata['high']= df.iloc[i,2]
        stockdata['low']= df.iloc[i,3]
        stockdata['close']= df.iloc[i,4]
        stockdata['volume']= df.iloc[i,5]
        if stockdata['volume'] =='':
           stockdata['volume'] ='0'
        
        stockdata['amount']= df.iloc[i,6]
        if stockdata['amount'] =='':
           stockdata['amount'] ='0'
        
        stockdata['turn']= df.iloc[i,7]    
        if stockdata['turn'] =='':
           stockdata['turn'] ='0'
        
        stockdata['pctChg']= df.iloc[i,8]
        if stockdata['pctChg'] =='':
           stockdata['pctChg'] ='0'
           
        #stockdata['float_share']= df.iloc[i,9]
        #if stockdata['float_share'] =='':
        #   stockdata['float_share'] ='0'
           
        #print(stockdata['float_share'])   
           
        
        # #查找是否有指定日期的股票数据
        result = mysql.find_stock_day_data(db,code_name,stockdata['date'])
        #没有则添加
        if result == 0:
            mysql.add_stock_day_data(db, code_name, stockdata)
            
                       
#主程序入口
if __name__ == '__main__':
    #download_create_all_database() 
    download_update_all_stock_day('2023-5-17','2023-5-18')
     
